Conversation
Replace calls to `_get_sp_group_from_device_mesh` with direct access to `sequence_parallel._sp_group` in sequence parallel attention tests. This simplifies the test setup by using the already initialized group stored in the module, improving code clarity and reducing redundancy.
…d functions Replace manual gradient handling with `torch.autograd.Function` subclasses `_ReduceSequenceParallelLoss` and `_ReduceSequenceParallelSum` to compute global loss via autograd-aware all-reduce. This simplifies the logic for both sum and mean reductions, improves gradient correctness, and removes the need for separate metric scaling when `world_size > 1`.
Add `compensate_fsdp_avg` config flag to adjust loss reduction when sequence parallel (SP) is combined with FSDP or accelerate DDP/FSDP. This prevents gradient magnitude from being incorrectly scaled down by an extra factor of SP world size during data-parallel averaging. - In `GatherLoss` backward, scale gradients by SP world size before splitting, so downstream FSDP averaging does not shrink this path. - In `SequenceParallelStrategy.reduce_loss`, apply a compensation factor (ulysses_size) when `compensate_fsdp_avg` is enabled. - Automatically set `compensate_fsdp_avg=True` in `TransformersModel` when using NativeFSDPStrategy or AccelerateStrategy with both SP and data parallelism active.
- Add 'kernels' as an optional dependency group in pyproject.toml - Refactor CI container test script to use a reusable installation function - Install twinkle with kernels in both debug and release modes for consistency - Improve maintainability by centralizing the installation command
Update `_load_from_hub` function to handle API changes in `select_revision_or_version` and `get_kernel` calls. The changes introduce try-except blocks to catch `TypeError` exceptions, allowing the function to work with both modern keyword-based APIs and older positional argument variants. This ensures compatibility across different versions of the kernels module without breaking existing functionality.
Add additional import statements and validation steps to ensure the required kernel 'kernels-test/flattened-build' can be successfully loaded before proceeding with the test. This prevents test failures due to missing or incompatible kernel environments and provides clearer skip messages when the kernel is unavailable.
Summary of ChangesHello @meichangsu1, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several important changes, primarily focused on improving sequence parallel training. Key changes include adding kernels as an optional dependency, ensuring backward compatibility with different kernels library API versions, and implementing significant correctness fixes for loss calculation and gradient scaling in sequence parallelism, especially when combined with data parallelism. A bug in label masking for packed sequences is also fixed. My review includes a suggestion to improve code structure by refactoring nested classes for better maintainability. The pull request title 'skip unitest' is not very descriptive of these substantial changes; a more informative title would be helpful.
| class _ReduceSequenceParallelLoss(torch.autograd.Function): | ||
|
|
||
| @staticmethod | ||
| def forward(ctx, local_mean: torch.Tensor, num_valid_tokens: torch.Tensor) -> torch.Tensor: | ||
| local_tokens = num_valid_tokens.detach().clone() | ||
| local_sum = local_mean * local_tokens | ||
| if local_tokens.item() == 0: | ||
| local_sum = torch.nan_to_num(local_sum) | ||
| global_sum = local_sum.detach().clone() | ||
| dist.all_reduce(global_sum, group=sequence_parallel._sp_group) | ||
| global_tokens = num_valid_tokens.detach().clone() | ||
| dist.all_reduce(global_tokens, group=sequence_parallel._sp_group) | ||
| ctx.save_for_backward(local_tokens, global_tokens) | ||
| if global_tokens.item() == 0: | ||
| return local_sum | ||
| return global_sum / global_tokens | ||
|
|
||
| @staticmethod | ||
| def backward(ctx, grad_output: torch.Tensor): | ||
| local_tokens, global_tokens = ctx.saved_tensors | ||
| if global_tokens.item() == 0: | ||
| return torch.zeros_like(grad_output), None | ||
| # d(global_mean)/d(local_mean) = local_tokens / global_tokens. | ||
| grad_local_mean = grad_output * (local_tokens / global_tokens) * compensate_factor | ||
| return grad_local_mean, None | ||
|
|
||
| class _ReduceSequenceParallelSum(torch.autograd.Function): | ||
|
|
||
| @staticmethod | ||
| def forward(ctx, local_sum: torch.Tensor) -> torch.Tensor: | ||
| ctx.sum_metric_scale = sum_metric_scale | ||
| global_sum = local_sum.detach().clone() | ||
| dist.all_reduce(global_sum, group=sequence_parallel._sp_group) | ||
| # Keep logging/metric value aligned with non-SP sum semantics under | ||
| # outer collect='mean' by removing one SP replication factor. | ||
| return global_sum / ctx.sum_metric_scale | ||
|
|
||
| @staticmethod | ||
| def backward(ctx, grad_output: torch.Tensor): | ||
| # Keep training gradient scale unchanged; forward-side scaling is for | ||
| # logging/metric alignment under outer collect='mean'. | ||
| return grad_output | ||
|
|
There was a problem hiding this comment.
For better code structure and maintainability, consider moving the nested classes _ReduceSequenceParallelLoss and _ReduceSequenceParallelSum to the module level. Defining them outside the reduce_loss method would make them reusable, easier to test in isolation, and improve the readability of the reduce_loss method itself. You would need to pass context variables like compensate_factor, sum_metric_scale, and sequence_parallel._sp_group as arguments to the apply method.
2 participants
No description provided.